# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

import __main__

from coin.base.datetime_util import to_datetime
from coin.base.hostname import get_hostname
from coin.base.timestamp import get_timestamp
from coin.exchange.base.proto_logger import (
    ProtoLogProvider,
    to_date_str,
    to_time_str,
)
from coin.proto.coin_request_pb2 import (
    AccountRequestProto,
    StrategyRequestProto,
)
from coin.proto.coin_log_pb2 import (
    LogFileResponseProto,)
from coin.exchange.base.order_gateway_log_builder import (prepare_builder,
                                                          to_proto_order_side,
                                                          to_proto_order_type,
                                                          BaseOrderEventBuilder,
                                                          BaseRequestBuilder,
                                                          BaseResponseBuilder)
from coin.proto.coin_order_gateway_pb2 import (OrderAmendRequestProto,
                                               OrderAmendResponseProto,
                                               OrderCancelRequestProto,
                                               OrderCancelResponseProto,
                                               OrderEvent,
                                               OrderGatewayLog,
                                               OrderSubmitRequestProto,
                                               OrderSubmitResponseProto)
from coin.proto.coin_query_pb2 import (AccountBalance,
                                       AccountBalanceEstimate,
                                       AccountPosition,
                                       AccountPositionEstimate)
import coin.proto.coin_order_enums_pb2 as coin_order
from coin.flow.archive import ensure_dir_exists
from coin.flow.topic_record import TopicRecordWriter

default_log_root = os.path.expanduser('~/data/og_proto_log')


def gen_og_log_dir(*, log_root, trading_date, machine):
  date_str = trading_date.strftime('%Y%m%d')
  log_dir = os.path.join(log_root, machine, date_str)
  return log_dir


def gen_og_log_request(config, og):
  if config.owner is not None and config.key_name is not None:
      owner = config.key_name
  else:
      owner = config.owner
  return AccountRequestProto(market_type=og.market_type,
                             exchange=og.exchange,
                             api_version=og.api_version,
                             owner=owner)


def compose_file_name(request: AccountRequestProto, time_str: str, timestamp):
  assert type(request) == AccountRequestProto
  main_file = os.path.basename(__main__.__file__)
  log_name = '%s.%s.%s.%s.%s' % (
      request.owner, request.market_type, request.exchange, time_str, timestamp)
  return log_name


# <strategy_name>.py.<owner>.<market_type>.<exchange>.<YYmmDD-HHMMSS>.<timestamp_nano_sec>
# eg: test_mm_ma_strategy2.py.jshin0.Futures.Bitmex.20180611-084408.1528706648659898000
def parse_file_name(file_name):
  exts = ('.gz', '.lz4')
  for ext in exts:
    if file_name.endswith(ext):
      file_name = file_name[:-len(ext)]
      break

  exploded = file_name.split('.')
  if len(exploded) == 6:
    # This is for old format which does not have timestamp at the end.
    # <strategy_name>.py.<owner>.<market_type>.<exchange>.<YYmmDD-HHMMSS>
    timestamp = None
  elif len(exploded) == 7:
    timestamp = exploded[-1]
    exploded = exploded[:-1]
  elif len(exploded) == 8:
    # This is the case when synced log file is incomplete. There will be a
    # random suffix. Invalid and should be skipped.
    return None
  else:
    raise ValueError('Wrong file name: %s' % file_name)

  time_str = exploded[-1]
  exchange = exploded[-2]
  market_type = exploded[-3]
  owner = exploded[-4]
  # exploded[-5] is `.py`
  assert exploded[-5] == 'py'
  strategy_group = exploded[-6]
  _date, _time = time_str.split('-')
  acct_req = AccountRequestProto(market_type=market_type, exchange=exchange, owner=owner)
  strat_req = StrategyRequestProto(strategy_group=strategy_group)
  return LogFileResponseProto(type=LogFileResponseProto.OG_LOG,
                              log_file_name=file_name,
                              date=_date,
                              time=_time,
                              timestamp=timestamp,
                              acct_req=acct_req,
                              strat_req=strat_req)


# <owner>.<market_type>.<exchange>.<YYmmDD-HHMMSS>.<timestamp_nano_sec>.gz
# eg: prestoinvt67.Futures.Huobi.20200222-182610.1582395970456699143.gz
def parse_file_name_from_coin2(file_name):
  exts = ('.gz', '.lz4')
  for ext in exts:
    if file_name.endswith(ext):
      file_name = file_name[:-len(ext)]
      break

  exploded = file_name.split('.')
  if len(exploded) == 5:
    timestamp = exploded[-1]
    exploded = exploded[:-1]
  elif len(exploded) == 8:
    # This is the case when synced log file is incomplete. There will be a
    # random suffix. Invalid and should be skipped.
    return None
  else:
    raise ValueError('Wrong file name: %s' % file_name)

  time_str = exploded[-1]
  exchange = exploded[-2]
  market_type = exploded[-3]
  owner = exploded[-4]
  _date, _time = time_str.split('-')
  acct_req = AccountRequestProto(market_type=market_type, exchange=exchange, owner=owner)
  strat_req = StrategyRequestProto()
  return LogFileResponseProto(type=LogFileResponseProto.OG_LOG,
                              log_file_name=file_name,
                              date=_date,
                              time=_time,
                              timestamp=timestamp,
                              acct_req=acct_req,
                              strat_req=strat_req)


def enumerate_logs(request: AccountRequestProto, trading_date, log_dir):
  assert type(request) == AccountRequestProto
  trading_date = trading_date.strftime('%Y%m%d')
  files_dict = {}
  if not os.path.isdir(log_dir):
    return []
  for i, file_name in enumerate(os.listdir(log_dir)):
    try:
      resp = parse_file_name(file_name)
    except Exception:
      resp = parse_file_name_from_coin2(file_name)
    if resp is None:
      continue
    if (resp.acct_req.market_type == request.market_type
        and resp.acct_req.exchange == request.exchange and resp.date == trading_date):
      if request.HasField('owner') and resp.acct_req.owner != request.owner:
        continue
      _time = int(resp.time)
      files_dict[(_time, i)] = os.path.join(log_dir, file_name)

  return [files_dict[key] for key in sorted(files_dict.keys())]


class OrderGatewayLogProvider(ProtoLogProvider):
  def __init__(self, request: AccountRequestProto, log_root=None, machine=None):
    assert type(request) == AccountRequestProto
    self._request = request
    self._log_root = log_root or default_log_root
    self._machine = machine or get_hostname()
    self._writer = None
    super().__init__()

    self._order_event_id = 0

  @property
  def request(self):
    return self._request

  def get_next_order_event_id(self):
    self._order_event_id += 1
    return self._order_event_id

  def _gen_log_path_and_ensure(self, timestamp):
    log_root = self._log_root
    machine = self._machine
    trading_date = to_datetime(timestamp)
    date_str = trading_date.strftime('%Y%m%d')
    time_str = to_time_str(timestamp)
    log_name = compose_file_name(request=self.request, time_str=time_str, timestamp=timestamp)

    log_dir = gen_og_log_dir(
        log_root=log_root, trading_date=trading_date, machine=machine)
    ensure_dir_exists(log_dir)
    log_path = os.path.join(log_dir, log_name)

    return LogFileResponseProto(type=LogFileResponseProto.OG_LOG,
                                log_file_name=log_path,
                                date=date_str,
                                time=time_str,
                                timestamp=str(timestamp),
                                acct_req=self.request)

  def gen_writer(self, timestamp):
    log_file_name = self._gen_log_path_and_ensure(timestamp).log_file_name
    if self._writer is None or self._writer.filepath != log_file_name:
      self._writer = TopicRecordWriter(log_file_name)
    return self._writer

  def write_account_request(self, flush: bool = True):
    log_proto = OrderGatewayLog(type=OrderGatewayLog.ACCOUNT_REQUEST, account_request=self.request)
    self.write(log_proto, flush)

  def write_position(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, AccountPosition)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.POSITION,
        position=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_balance(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, AccountBalance)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.BALANCE,
        balance=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_position_estimate(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, AccountPositionEstimate)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.POSITION_ESTIMATE,
        position_estimate=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_balance_estimate(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, AccountBalanceEstimate)
    log_proto = OrderGatewayLog(
      type=OrderGatewayLog.BALANCE_ESTIMATE,
      balance_estimate=protobuf,
      account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_private_exchange_message(self, private_exchange_message, flush: bool = True):
    assert isinstance(private_exchange_message, str)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.PRIVATE_EXCHANGE_MESSAGE,
        private_exchange_message=private_exchange_message,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_event(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderEvent)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.ORDER_EVENT,
        event=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_submit_request(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderSubmitRequestProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.SUBMIT_REQUEST,
        submit_request=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_submit_response(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderSubmitResponseProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.SUBMIT_RESPONSE,
        submit_response=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_amend_request(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderAmendRequestProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.AMEND_REQUEST,
        amend_request=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_amend_response(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderAmendResponseProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.AMEND_RESPONSE,
        amend_response=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_cancel_request(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderCancelRequestProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.CANCEL_REQUEST,
        cancel_request=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def write_order_cancel_response(self, protobuf, flush: bool = True):
    assert isinstance(protobuf, OrderCancelResponseProto)
    log_proto = OrderGatewayLog(
        type=OrderGatewayLog.CANCEL_RESPONSE,
        cancel_response=protobuf,
        account_request=self.request,
    )
    self.write(log_proto, flush)

  def gen_order_request_and_order_event_and_log(self,
                                                symbol,
                                                order_price,
                                                order_qty,
                                                order_side,
                                                order_type,
                                                internal_order_id,
                                                timestamp):
    req_builder = prepare_builder(BaseRequestBuilder,
                                  OrderSubmitRequestProto,
                                  timestamp,
                                  self.get_next_request_id(),
                                  self.exchange_type,
                                  symbol)
    req_builder.request.internal_order_id = internal_order_id
    self.write_order_submit_request(req_builder.request)

    oe_builder = prepare_builder(BaseOrderEventBuilder,
                                 None,
                                 timestamp,
                                 self.get_next_order_event_id(),
                                 self.exchange_type,
                                 symbol)

    order_side = to_proto_order_side(order_side)
    order_type = to_proto_order_type(order_type)
    oe_builder.set_order_info(internal_order_id, order_price, order_qty, order_side, order_type)

    oe_builder.set_state(order_event_type=OrderEvent.ORDER_SUBMITTED,
                         order_state=coin_order.PENDING_ORDER,
                         order_event_source=OrderEvent.USER)
    self.write_order_event(oe_builder.oe)

  def gen_order_response_and_order_event_and_log(self,
                                                 symbol,
                                                 success,
                                                 internal_order_id,
                                                 timestamp):
    resp_builder = prepare_builder(BaseResponseBuilder,
                                   OrderSubmitResponseProto,
                                   timestamp,
                                   self.get_next_response_id(),
                                   self.exchange_type,
                                   symbol)

    resp_builder.response.success = success
    resp_builder.response.internal_order_id = internal_order_id
    self.write_order_submit_response(resp_builder.response)

    if not success:
      oe_builder = prepare_builder(BaseOrderEventBuilder,
                                   None,
                                   timestamp,
                                   self.get_next_order_event_id(),
                                   self.exchange_type,
                                   symbol)
      oe_builder.oe.internal_order_id = internal_order_id

      oe_builder.set_state(order_event_type=OrderEvent.ORDER_REJECTED,
                           order_state=coin_order.DEAD_ORDER,
                           order_event_source=OrderEvent.EXCHANGE)
      self.write_order_event(oe_builder.oe)

  def gen_cancel_request_and_order_event_and_log(self, symbol, internal_order_id, timestamp):
    req_builder = prepare_builder(BaseRequestBuilder,
                                  OrderCancelRequestProto,
                                  timestamp,
                                  self.get_next_request_id(),
                                  self.exchange_type,
                                  symbol)
    if internal_order_id:
      req_builder.request.internal_order_id = internal_order_id
    self.write_order_cancel_request(req_builder.request)

    oe_builder = prepare_builder(BaseOrderEventBuilder,
                                 None,
                                 timestamp,
                                 self.get_next_order_event_id(),
                                 self.exchange_type,
                                 symbol)
    if internal_order_id:
      oe_builder.oe.internal_order_id = internal_order_id

    oe_builder.set_state(order_event_type=OrderEvent.CANCEL_SUBMITTED,
                         order_state=coin_order.WORKING_ORDER,
                         order_event_source=OrderEvent.USER)
    self.write_order_event(oe_builder.oe)

  def gen_cancel_response_and_log(self, symbol, success, internal_order_id):
    resp_builder = prepare_builder(BaseResponseBuilder,
                                   OrderCancelResponseProto,
                                   None,
                                   self.get_next_response_id(),
                                   self.exchange_type,
                                   symbol)
    resp_builder.response.success = success
    if internal_order_id:
      resp_builder.response.internal_order_id = internal_order_id
    self.write_order_cancel_response(resp_builder.response)

  def gen_flush_unacked_order_event_and_log(self, symbol, internal_order_id, timestamp):
    builder = prepare_builder(BaseOrderEventBuilder,
                              None,
                              timestamp,
                              self.get_next_order_event_id(),
                              self.exchange_type,
                              symbol)
    if internal_order_id is not None:
      builder.oe.internal_order_id = internal_order_id

    builder.set_state(order_event_type=OrderEvent.ORDER_REJECTED,
                      order_state=coin_order.DEAD_ORDER,
                      order_event_source=OrderEvent.USER)
    self.write_order_event(builder.oe)
